﻿/*
 * CAudioEncoder.cpp
 *
 *  Created on: 2017年9月29日
 *      Author: terry
 */

#include "CAudioEncoder.h"
#include "AVFramePtr.h"
#include "DateTime.h"
#include "CLog.h"
#include "AacHelper.h"
#include "CodecIDConvert.h"


using namespace av;


CAudioEncoder::CAudioEncoder():
	m_format(),
	m_context(),
	m_swrContext(),
	m_fifo(),
	m_sampleCount(),
	m_isAdts(false),
	m_frameSize()
{
	m_buffer.ensure(1024 * 48);
}

CAudioEncoder::~CAudioEncoder()
{
	close();
}

int CAudioEncoder::open(int channels, int rate, int codec, int bitrate)
{
	CLog::debug("CAudioEncoder::open. channels:%d, rate:%d, bitrate:%d, codec:%d\n",
			channels, rate, bitrate, codec);

	m_format.audioCodec = codec;
	m_format.channels = channels;
	m_format.sampleRate = rate;
	m_format.audioBitrate = bitrate;
	m_format.audioRate = rate;
	m_format.audioProfile = FF_PROFILE_AAC_LOW;

	return openEncoder();
}

void CAudioEncoder::close()
{
	closeEncoder();

	closeFifo();

    m_sampleCount = 0;
}

bool CAudioEncoder::isOpen()
{
	return (m_context != NULL);
}

int CAudioEncoder::encode(uint8_t* data, int size, int64_t pts, MPacket* pkt)
{
	if (getFifoSize() >= m_frameSize)
	{
		encodeFrame(pkt);
		return 0;
	}

	AVFramePtr outFrame(av_frame_alloc(), avframe_free);
	outFrame->channels = m_format.channels;
	outFrame->sample_rate = m_format.sampleRate;
	outFrame->channel_layout = av_get_default_channel_layout(outFrame->channels);
	outFrame->format = AV_SAMPLE_FMT_S16;
	outFrame->pts = pts;
	outFrame->nb_samples = size / 2 / m_format.channels;

	av_frame_get_buffer(outFrame.get(), 1);
	memcpy(outFrame->data[0], data, size);
	outFrame->linesize[0] = size;

	AVPacket outPkt;
	av_init_packet(&outPkt);
	outPkt.data = m_buffer.data() + AacHelper::ADTS_HEADER_SIZE;
	outPkt.size = m_buffer.capacity() - AacHelper::ADTS_HEADER_SIZE;

	if (encodeFrame(outFrame.get(), &outPkt))
	{
		makePacket(outPkt, pkt);
	}

	return size;
}

bool CAudioEncoder::encodeFrame(AVFrame* frame, AVPacket* pkt)
{
	int rc = 0;
	if (!isEncoderOpen())
	{
		m_format.channels = frame->channels;
		m_format.sampleRate = frame->sample_rate;

		rc = openEncoder();
		if (rc != 0)
		{
			return false;
		}
	}

	if (frame)
	{
		if (isSamleFormat(frame->format, frame->channels, frame->sample_rate))
		{
            storeFrame(frame);
        }
        else
        {
					/// resample
            AVFramePtr outFrame(av_frame_alloc(), avframe_free);

            outFrame->format = m_context->sample_fmt;
            outFrame->channels = m_context->channels;
            outFrame->channel_layout = m_context->channel_layout;
            outFrame->sample_rate = m_context->sample_rate;
            outFrame->pts = frame->pts;

			rc = resample(frame, outFrame.get());
			if (rc > 0)
			{
				storeFrame(outFrame.get());
			}
			else
			{
				//return rc;
			}
		}
	}

	if (getFifoSize() < m_context->frame_size)
	{
		return false;
	}

	return encodeFrame(pkt);

}

void CAudioEncoder::flush()
{
	if (m_fifo)
	{
		av_audio_fifo_reset(m_fifo);
	}
}

MFormat& CAudioEncoder::getFormat()
{
	return m_format;
}

bool CAudioEncoder::setParam(int param, int value)
{
	if (param == kAdts)
	{
		m_isAdts = (value > 0);
	}
	return true;
}

void CAudioEncoder::setAdts(bool isAdts)
{
	m_isAdts = isAdts;
}

int CAudioEncoder::openEncoder()
{
	//AV_CODEC_ID_AAC;
	AVCodecID codecID = CodecIDConvert::toAVCodecID(m_format.audioCodec);
	AVCodec* codec = avcodec_find_encoder(codecID);
	if (!codec)
	{
		CLog::warning("no such codec. %d\n", codecID);
		return ENOENT;
	}

	m_context = avcodec_alloc_context3(codec);
	if (!m_context)
	{
		return ENOENT;
	}

    if (codecID == AV_CODEC_ID_AAC)
    {
	    m_context->profile = m_format.audioProfile;
    }

	m_context->channels = m_format.channels;
    m_context->channel_layout = av_get_default_channel_layout(m_format.channels);
	m_context->sample_rate = m_format.sampleRate;
    m_context->sample_fmt = codec->sample_fmts[0];
    m_context->time_base = av_make_q(1, m_format.sampleRate);

    AVDictionary* options = NULL;
    if (codecID == AV_CODEC_ID_AAC)
    {
        av_dict_set(&options, "strict", "-2", 0);
    }

    int rc = avcodec_open2(m_context, codec, &options);

    av_dict_free(&options);

    if (rc < 0)
    {
        char buf[256];
        av_strerror(rc, buf, 256);
		
		CLog::warning("failed to open codec. rc:%d, %s\n", rc, buf);

        av_free(m_context);
        m_context = NULL;

        return ENOENT;
    }

    if (m_context->extradata_size > 0)
    {
        m_extradata.assign((char*)m_context->extradata, m_context->extradata_size);
        m_format.setConfig(m_context->extradata, m_context->extradata_size);
    }

	m_frameSize = m_context->frame_size;

	return 0;
}

void CAudioEncoder::closeEncoder()
{
	if (m_swrContext)
	{
		swr_close(m_swrContext);
		swr_free(&m_swrContext);
		m_swrContext = NULL;
	}

	if (m_context)
	{
		avcodec_close(m_context);
		av_free(m_context);
		m_context = NULL;
	}
}

bool CAudioEncoder::isEncoderOpen()
{
	return (m_context != NULL);
}

bool CAudioEncoder::isSamleFormat(int fmt, int channels, int sampleRate)
{
	return (m_format.channels == channels) &&
        (m_format.sampleRate == sampleRate) &&
        (fmt == m_context->sample_fmt);
}

bool CAudioEncoder::getAudioConfig(std::string& config)
{
    config = m_extradata;
	return true;
}

int CAudioEncoder::resample(AVFrame* inFrame, AVFrame* outFrame)
{
	if (m_swrContext == NULL)
	{
		m_swrContext = swr_alloc();

		swr_config_frame(m_swrContext, outFrame, inFrame);

		int rc = swr_init(m_swrContext);
		if (rc != 0)
		{
			CLog::error("swr_init failed. rc:%d\n", rc);
		}
	}
	else
	{
		int delay = swr_get_delay(m_swrContext, outFrame->sample_rate);

		outFrame->nb_samples = inFrame->nb_samples + delay;
		av_frame_get_buffer(outFrame, 1);
	}

	swr_convert_frame(m_swrContext, outFrame, inFrame);

	return outFrame->nb_samples;
}

void CAudioEncoder::storeFrame(AVFrame* frame)
{
	if (!m_fifo)
	{
		m_fifo = av_audio_fifo_alloc((AVSampleFormat)frame->format, frame->channels, frame->nb_samples);
	}

    av_audio_fifo_write(m_fifo, (void**)frame->data, frame->nb_samples);
}

bool CAudioEncoder::loadFrame(AVFrame* frame, int nb_samples)
{
	if (!m_fifo)
	{
		return false;
	}

	if (av_audio_fifo_size(m_fifo) < nb_samples)
	{
		return false;
	}

	av_audio_fifo_read(m_fifo, (void**)frame->data, nb_samples);
	return true;
}

void CAudioEncoder::closeFifo()
{
	if (m_fifo)
	{
		av_audio_fifo_free(m_fifo);
		m_fifo = NULL;
	}
}

int CAudioEncoder::getFifoSize()
{
	if (!m_fifo)
	{
		return 0;
	}
	return av_audio_fifo_size(m_fifo);
}


bool CAudioEncoder::encodeFrame(AVPacket* pkt)
{
	AVFramePtr encFrame(av_frame_alloc(), avframe_free);
	AVFrame* frame = encFrame.get();
	frame->nb_samples = m_context->frame_size;
	frame->channels = m_context->channels;
	frame->channel_layout = m_context->channel_layout;
	frame->sample_rate = m_context->sample_rate;
	frame->format = m_context->sample_fmt;

	av_frame_get_buffer(frame, 0);

	if (!loadFrame(encFrame.get(), m_context->frame_size))
	{
		return false;
	}

	m_sampleCount += m_context->frame_size;
	frame->pts = av_rescale(m_sampleCount, AV_TIME_BASE, frame->sample_rate);

	int got_packet = 0;
	int rc = avcodec_encode_audio2(m_context, pkt, frame, &got_packet);
	if (rc == 0)
	{
		if (got_packet)
		{
			pkt->duration = frame->nb_samples;
			pkt->pos = comn::DateTime::now().totalMillisecond();
		}
	}

	return got_packet > 0;
}

bool CAudioEncoder::encodeFrame(MPacket* pkt)
{
	bool done = false;
	AVPacket outPkt;
	av_init_packet(&outPkt);
	outPkt.data = m_buffer.data() + AacHelper::ADTS_HEADER_SIZE;
	outPkt.size = m_buffer.capacity() - AacHelper::ADTS_HEADER_SIZE;

	if (encodeFrame(&outPkt))
	{
		makePacket(outPkt, pkt);

		done = true;
	}
	return done;
}

void CAudioEncoder::makePacket(AVPacket& avpkt, MPacket* pkt)
{
	if (m_isAdts)
	{
		int profile = FF_PROFILE_AAC_LOW;
		int packetLen = avpkt.size;
		av::AacHelper::makeAdtsHeader(profile, m_format.sampleRate, m_format.channels,
			packetLen + av::AacHelper::ADTS_HEADER_SIZE,
			m_buffer.data(), av::AacHelper::ADTS_HEADER_SIZE);

		avpkt.data -= av::AacHelper::ADTS_HEADER_SIZE;
		avpkt.size += av::AacHelper::ADTS_HEADER_SIZE;
	}

	pkt->type = MTYPE_AUDIO;
	pkt->data = avpkt.data;
	pkt->size = avpkt.size;
	pkt->flags = avpkt.flags;
	pkt->duration = avpkt.duration;
	pkt->pts = avpkt.pts;
}